import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score
from tensorflow.keras.datasets import imdb

# 加载IMDB数据集
max_features = 5000  # 词汇表大小
(X_train, y_train), (X_test, y_test) = imdb.load_data(num_words=max_features)

# 获取单词索引字典
word_index = imdb.get_word_index()
reverse_word_index = {v: k for k, v in word_index.items()}

# 将整数序列转换回文本
def decode_review(sequence):
    return ' '.join([reverse_word_index.get(i - 3, '?') for i in sequence])

X_train_text = [decode_review(seq) for seq in X_train]
X_test_text = [decode_review(seq) for seq in X_test]

# 使用TF-IDF进行特征提取
vectorizer = TfidfVectorizer(max_features=max_features)
X_train_tfidf = vectorizer.fit_transform(X_train_text)
X_test_tfidf = vectorizer.transform(X_test_text)

# 逻辑回归模型
lr_model = LogisticRegression()
lr_model.fit(X_train_tfidf, y_train)
y_pred_lr = lr_model.predict(X_test_tfidf)
lr_accuracy = accuracy_score(y_test, y_pred_lr)
print(f"Logistic Regression Accuracy: {lr_accuracy}")

# SVM模型
svm_model = SVC(kernel='linear')  # 使用线性核
svm_model.fit(X_train_tfidf, y_train)
y_pred_svm = svm_model.predict(X_test_tfidf)
svm_accuracy = accuracy_score(y_test, y_pred_svm)
print(f"SVM Accuracy: {svm_accuracy}")

# # 比较准确率差异
# print("\nAccuracy Comparison:")
# print(f"Logistic Regression: {lr_accuracy:.4f}")
# print(f"SVM: {svm_accuracy:.4f}")